package io.quarkus.arc.processor;

import static org.jboss.jandex.gizmo2.Jandex2Gizmo.classDescOf;

import java.lang.constant.ClassDesc;
import java.lang.constant.ConstantDescs;
import java.util.ArrayList;
import java.util.List;
import java.util.Objects;
import java.util.function.Predicate;

import org.jboss.jandex.AnnotationInstance;
import org.jboss.jandex.AnnotationValue;
import org.jboss.jandex.ArrayType;
import org.jboss.jandex.ClassInfo;
import org.jboss.jandex.DotName;
import org.jboss.jandex.IndexView;
import org.jboss.jandex.MethodInfo;
import org.jboss.jandex.Type;

import io.quarkus.arc.impl.ComputingCache;
import io.quarkus.gizmo2.Const;
import io.quarkus.gizmo2.Expr;
import io.quarkus.gizmo2.LocalVar;
import io.quarkus.gizmo2.creator.BlockCreator;
import io.quarkus.gizmo2.desc.ConstructorDesc;
import io.quarkus.gizmo2.desc.FieldDesc;

/**
 * Handles generating bytecode for annotation literals. The
 * {@link #create(BlockCreator, ClassInfo, AnnotationInstance) create()} method can be used
 * to generate a bytecode sequence for instantiating an annotation literal.
 * <p>
 * Behind the scenes, for each annotation literal, its class is also generated. This class
 * is supposed to be used at runtime. The generated annotation literal classes are shared.
 * That is, one class is generated for each annotation type and the constructor of that class
 * accepts values of all annotation members. As a special case, annotation literal classes
 * for memberless annotation types have a singleton instance.
 * <p>
 * This construct is thread-safe.
 */
public class AnnotationLiteralProcessor {
    private static final String ANNOTATION_LITERAL_SUFFIX = "_ArcAnnotationLiteral";

    private final ComputingCache<CacheKey, AnnotationLiteralClassInfo> cache;
    private final IndexView beanArchiveIndex;

    AnnotationLiteralProcessor(IndexView beanArchiveIndex, Predicate<DotName> applicationClassPredicate) {
        this.cache = new ComputingCache<>(key -> new AnnotationLiteralClassInfo(
                generateAnnotationLiteralClassName(key.annotationName()),
                applicationClassPredicate.test(key.annotationName()),
                key.annotationClass));
        this.beanArchiveIndex = Objects.requireNonNull(beanArchiveIndex);
    }

    boolean hasLiteralsToGenerate() {
        return !cache.isEmpty();
    }

    ComputingCache<CacheKey, AnnotationLiteralClassInfo> getCache() {
        return cache;
    }

    /**
     * Generates a bytecode sequence to create an instance of given annotation type, such that
     * the annotation members have the same values as the given annotation instance.
     * An implementation of the annotation type will be generated automatically.
     * <p>
     * It is expected that given annotation instance is runtime-retained; an exception is thrown
     * if not. Further, it is expected that the annotation type is available (that is,
     * {@code annotationClass != null}); an exception is thrown if not. Callers that expect
     * they always deal with runtime-retained annotations whose classes are available do not
     * have to check (and will get decent errors for free), but callers that can possibly deal
     * with class-retained annotations or missing annotation classes must check explicitly.
     * <p>
     * We call the generated implementation of the annotation type an <em>annotation literal class</em>
     * and the instance produced by the generated bytecode an <em>annotation literal instance</em>,
     * even though the generated code doesn't use CDI's {@code AnnotationLiteral}.
     *
     * @param bc will receive the bytecode sequence for instantiating the annotation literal class
     *        as a sequence of {@link BlockCreator} method calls
     * @param annotationClass the annotation type
     * @param annotationInstance the annotation instance; must match the {@code annotationClass}
     * @return an annotation literal instance result handle
     */
    public Expr create(BlockCreator bc, ClassInfo annotationClass, AnnotationInstance annotationInstance) {
        if (!annotationInstance.runtimeVisible()) {
            throw new IllegalArgumentException("Annotation does not have @Retention(RUNTIME): " + annotationInstance);
        }
        if (annotationClass == null) {
            throw new IllegalArgumentException("Annotation class not available: " + annotationInstance);
        }

        AnnotationLiteralClassInfo literal = cache.getValue(new CacheKey(annotationClass));

        ClassDesc generatedClass = ClassDesc.of(literal.generatedClassName);

        if (literal.annotationMembers().isEmpty()) {
            return bc.getStaticField(FieldDesc.of(generatedClass, "INSTANCE", generatedClass));
        }

        Expr[] ctorArgs = new Expr[literal.annotationMembers().size()];
        int argIndex = 0;
        for (MethodInfo annotationMember : literal.annotationMembers()) {
            AnnotationValue value = annotationInstance.value(annotationMember.name());
            if (value == null) {
                value = annotationMember.defaultValue();
            }
            if (value == null) {
                throw new IllegalStateException(String.format(
                        "Value is not set for %s.%s(). Most probably an older version of Jandex was used to index an application dependency. Make sure that Jandex 2.1+ is used.",
                        annotationMember.declaringClass().name(), annotationMember.name()));
            }
            Expr valueExpr = loadValue(bc, literal, annotationMember, value);
            ctorArgs[argIndex] = valueExpr;

            argIndex++;
        }

        ConstructorDesc ctor = ConstructorDesc.of(generatedClass, literal.annotationMembers()
                .stream()
                .map(it -> classDescOf(it.returnType()))
                .toArray(ClassDesc[]::new));
        return bc.new_(ctor, ctorArgs);
    }

    /**
     * Generates a bytecode sequence to load given annotation member value.
     *
     * @param bc will receive the bytecode sequence for loading the annotation member value
     *        as a sequence of {@link BlockCreator} method calls
     * @param literal data about the annotation literal class currently being generated
     * @param annotationMember the annotation member whose value we're loading
     * @param annotationMemberValue the annotation member value we're loading
     * @return the loaded annotation member value
     */
    private Expr loadValue(BlockCreator bc, AnnotationLiteralClassInfo literal,
            MethodInfo annotationMember, AnnotationValue annotationMemberValue) {

        return switch (annotationMemberValue.kind()) {
            case BOOLEAN -> Const.of(annotationMemberValue.asBoolean());
            case BYTE -> Const.of(annotationMemberValue.asByte());
            case SHORT -> Const.of(annotationMemberValue.asShort());
            case INTEGER -> Const.of(annotationMemberValue.asInt());
            case LONG -> Const.of(annotationMemberValue.asLong());
            case FLOAT -> Const.of(annotationMemberValue.asFloat());
            case DOUBLE -> Const.of(annotationMemberValue.asDouble());
            case CHARACTER -> Const.of(annotationMemberValue.asChar());
            case STRING -> Const.of(annotationMemberValue.asString());
            case ENUM -> {
                ClassDesc enumDesc = classDescOf(annotationMemberValue.asEnumType());
                yield bc.getStaticField(FieldDesc.of(enumDesc, annotationMemberValue.asEnum(), enumDesc));
            }
            case CLASS -> {
                if (annotationMemberValue.equals(annotationMember.defaultValue())) {
                    yield bc.getStaticField(FieldDesc.of(ClassDesc.of(literal.generatedClassName),
                            AnnotationLiteralGenerator.defaultValueStaticFieldName(annotationMember),
                            classDescOf(annotationMember.returnType())));
                } else {
                    yield Const.of(classDescOf(annotationMemberValue.asClass()));
                }
            }
            case NESTED -> {
                AnnotationInstance nestedAnnotation = annotationMemberValue.asNested();
                DotName annotationName = nestedAnnotation.name();
                ClassInfo annotationClass = beanArchiveIndex.getClassByName(annotationName);
                if (annotationClass == null) {
                    throw new IllegalStateException("Class of nested annotation " + nestedAnnotation + " missing");
                }
                yield create(bc, annotationClass, nestedAnnotation);
            }
            case ARRAY -> {
                AnnotationValue.Kind componentKind = annotationMemberValue.componentKind();
                yield switch (componentKind) {
                    case BOOLEAN -> {
                        boolean[] booleanArray = annotationMemberValue.asBooleanArray();
                        Expr[] exprArray = new Expr[booleanArray.length];
                        for (int i = 0; i < booleanArray.length; i++) {
                            exprArray[i] = Const.of(booleanArray[i]);
                        }
                        yield bc.newArray(boolean.class, exprArray);
                    }
                    case BYTE -> {
                        byte[] byteArray = annotationMemberValue.asByteArray();
                        Expr[] exprArray = new Expr[byteArray.length];
                        for (int i = 0; i < byteArray.length; i++) {
                            exprArray[i] = Const.of(byteArray[i]);
                        }
                        yield bc.newArray(byte.class, exprArray);
                    }
                    case SHORT -> {
                        short[] shortArray = annotationMemberValue.asShortArray();
                        Expr[] exprArray = new Expr[shortArray.length];
                        for (int i = 0; i < shortArray.length; i++) {
                            exprArray[i] = Const.of(shortArray[i]);
                        }
                        yield bc.newArray(short.class, exprArray);
                    }
                    case INTEGER -> {
                        int[] intArray = annotationMemberValue.asIntArray();
                        Expr[] exprArray = new Expr[intArray.length];
                        for (int i = 0; i < intArray.length; i++) {
                            exprArray[i] = Const.of(intArray[i]);
                        }
                        yield bc.newArray(int.class, exprArray);
                    }
                    case LONG -> {
                        long[] longArray = annotationMemberValue.asLongArray();
                        Expr[] exprArray = new Expr[longArray.length];
                        for (int i = 0; i < longArray.length; i++) {
                            exprArray[i] = Const.of(longArray[i]);
                        }
                        yield bc.newArray(long.class, exprArray);
                    }
                    case FLOAT -> {
                        float[] floatArray = annotationMemberValue.asFloatArray();
                        Expr[] exprArray = new Expr[floatArray.length];
                        for (int i = 0; i < floatArray.length; i++) {
                            exprArray[i] = Const.of(floatArray[i]);
                        }
                        yield bc.newArray(float.class, exprArray);
                    }
                    case DOUBLE -> {
                        double[] doubleArray = annotationMemberValue.asDoubleArray();
                        Expr[] exprArray = new Expr[doubleArray.length];
                        for (int i = 0; i < doubleArray.length; i++) {
                            exprArray[i] = Const.of(doubleArray[i]);
                        }
                        yield bc.newArray(double.class, exprArray);
                    }
                    case CHARACTER -> {
                        char[] charArray = annotationMemberValue.asCharArray();
                        Expr[] exprArray = new Expr[charArray.length];
                        for (int i = 0; i < charArray.length; i++) {
                            exprArray[i] = Const.of(charArray[i]);
                        }
                        yield bc.newArray(char.class, exprArray);
                    }
                    case STRING -> {
                        String[] stringArray = annotationMemberValue.asStringArray();
                        Expr[] exprArray = new Expr[stringArray.length];
                        for (int i = 0; i < stringArray.length; i++) {
                            exprArray[i] = Const.of(stringArray[i]);
                        }
                        yield bc.newArray(String.class, exprArray);
                    }
                    case ENUM -> {
                        String[] enumArray = annotationMemberValue.asEnumArray();
                        DotName[] enumTypeArray = annotationMemberValue.asEnumTypeArray();
                        LocalVar array = bc.localVar("array",
                                bc.newEmptyArray(componentTypeOf(annotationMember), enumArray.length));
                        for (int i = 0; i < enumArray.length; i++) {
                            ClassDesc classDesc = classDescOf(enumTypeArray[i]);
                            FieldDesc fieldDesc = FieldDesc.of(classDesc, enumArray[i], classDesc);
                            bc.set(array.elem(i), bc.getStaticField(fieldDesc));
                        }
                        yield array;
                    }
                    case CLASS -> {
                        if (annotationMemberValue.equals(annotationMember.defaultValue())) {
                            yield bc.getStaticField(FieldDesc.of(ClassDesc.of(literal.generatedClassName),
                                    AnnotationLiteralGenerator.defaultValueStaticFieldName(annotationMember),
                                    classDescOf(annotationMember.returnType())));
                        } else {
                            Type[] classArray = annotationMemberValue.asClassArray();
                            LocalVar array = bc.localVar("array",
                                    bc.newEmptyArray(componentTypeOf(annotationMember), classArray.length));
                            for (int i = 0; i < classArray.length; i++) {
                                bc.set(array.elem(i), Const.of(classDescOf(classArray[i])));
                            }
                            yield array;
                        }
                    }
                    case NESTED -> {
                        AnnotationInstance[] nestedArray = annotationMemberValue.asNestedArray();
                        LocalVar array = bc.localVar("array",
                                bc.newEmptyArray(componentTypeOf(annotationMember), nestedArray.length));
                        for (int i = 0; i < nestedArray.length; i++) {
                            AnnotationInstance nestedAnnotation = nestedArray[i];
                            DotName annotationName = nestedAnnotation.name();
                            ClassInfo annotationClass = beanArchiveIndex.getClassByName(annotationName);
                            if (annotationClass == null) {
                                throw new IllegalStateException("Class of nested annotation " + nestedAnnotation + " missing");
                            }
                            Expr nestedAnnotationValue = create(bc, annotationClass, nestedAnnotation);
                            bc.set(array.elem(i), nestedAnnotationValue);
                        }
                        yield array;
                    }
                    case UNKNOWN -> { // empty array
                        ClassDesc componentType = componentTypeOf(annotationMember);
                        // use empty array constants for common component kinds
                        if (ConstantDescs.CD_boolean.equals(componentType)) {
                            yield bc.getStaticField(FieldDescs.ANNOTATION_LITERALS_EMPTY_BOOLEAN_ARRAY);
                        } else if (ConstantDescs.CD_byte.equals(componentType)) {
                            yield bc.getStaticField(FieldDescs.ANNOTATION_LITERALS_EMPTY_BYTE_ARRAY);
                        } else if (ConstantDescs.CD_short.equals(componentType)) {
                            yield bc.getStaticField(FieldDescs.ANNOTATION_LITERALS_EMPTY_SHORT_ARRAY);
                        } else if (ConstantDescs.CD_int.equals(componentType)) {
                            yield bc.getStaticField(FieldDescs.ANNOTATION_LITERALS_EMPTY_INT_ARRAY);
                        } else if (ConstantDescs.CD_long.equals(componentType)) {
                            yield bc.getStaticField(FieldDescs.ANNOTATION_LITERALS_EMPTY_LONG_ARRAY);
                        } else if (ConstantDescs.CD_float.equals(componentType)) {
                            yield bc.getStaticField(FieldDescs.ANNOTATION_LITERALS_EMPTY_FLOAT_ARRAY);
                        } else if (ConstantDescs.CD_double.equals(componentType)) {
                            yield bc.getStaticField(FieldDescs.ANNOTATION_LITERALS_EMPTY_DOUBLE_ARRAY);
                        } else if (ConstantDescs.CD_char.equals(componentType)) {
                            yield bc.getStaticField(FieldDescs.ANNOTATION_LITERALS_EMPTY_CHAR_ARRAY);
                        } else if (ConstantDescs.CD_String.equals(componentType)) {
                            yield bc.getStaticField(FieldDescs.ANNOTATION_LITERALS_EMPTY_STRING_ARRAY);
                        } else if (ConstantDescs.CD_Class.equals(componentType)) {
                            yield bc.getStaticField(FieldDescs.ANNOTATION_LITERALS_EMPTY_CLASS_ARRAY);
                        } else {
                            yield bc.newEmptyArray(componentType, Const.of(0));
                        }
                    }
                    default -> {
                        // at this point, the only possible component kind is "array"
                        throw new UnsupportedOperationException("Array component kind is " + componentKind
                                + ", this should never happen");
                    }
                };
            }
            default -> throw new UnsupportedOperationException("Unsupported value: " + annotationMemberValue);
        };
    }

    private static ClassDesc componentTypeOf(MethodInfo annotationMember) {
        assert annotationMember.returnType().kind() == Type.Kind.ARRAY;
        return classDescOf(annotationMember.returnType().asArrayType().componentType());
    }

    // ---

    private static String componentType(MethodInfo method) {
        return componentTypeName(method).toString();
    }

    private static DotName componentTypeName(MethodInfo method) {
        ArrayType arrayType = method.returnType().asArrayType();
        return arrayType.constituent().name();
    }

    private static String generateAnnotationLiteralClassName(DotName annotationName) {
        // when the annotation is a java.lang annotation we need to use a different package in which to generate the literal
        // otherwise a security exception will be thrown when the literal is loaded
        boolean isJavaLang = annotationName.toString().startsWith("java.lang");
        String nameToUse = isJavaLang
                ? AbstractGenerator.DEFAULT_PACKAGE + annotationName.withoutPackagePrefix()
                : annotationName.toString();

        // com.foo.MyQualifier -> com.foo.MyQualifier_AnnotationLiteral
        return nameToUse + ANNOTATION_LITERAL_SUFFIX;
    }

    static class CacheKey {
        final ClassInfo annotationClass;

        CacheKey(ClassInfo annotationClass) {
            this.annotationClass = annotationClass;
        }

        DotName annotationName() {
            return annotationClass.name();
        }

        @Override
        public boolean equals(Object o) {
            if (this == o)
                return true;
            if (o == null || getClass() != o.getClass())
                return false;
            CacheKey cacheKey = (CacheKey) o;
            return Objects.equals(annotationClass.name(), cacheKey.annotationClass.name());
        }

        @Override
        public int hashCode() {
            return Objects.hashCode(annotationClass.name());
        }
    }

    static class AnnotationLiteralClassInfo {
        /**
         * Name of the generated annotation literal class.
         */
        final String generatedClassName;

        /**
         * Whether the generated annotation literal class is an application class.
         */
        final boolean isApplicationClass;

        /**
         * The annotation type. The generated annotation literal class will implement this interface.
         * The process that generates the annotation literal class may consult this, for example,
         * to know the set of annotation members.
         */
        final ClassInfo annotationClass;

        AnnotationLiteralClassInfo(String generatedClassName, boolean isApplicationClass, ClassInfo annotationClass) {
            this.generatedClassName = generatedClassName;
            this.isApplicationClass = isApplicationClass;
            this.annotationClass = annotationClass;
        }

        DotName annotationName() {
            return annotationClass.name();
        }

        List<MethodInfo> annotationMembers() {
            List<MethodInfo> result = new ArrayList<>();
            for (MethodInfo method : annotationClass.unsortedMethods()) {
                if (!method.name().equals(Methods.CLINIT) && !method.name().equals(Methods.INIT)) {
                    result.add(method);
                }
            }
            return result;
        }
    }
}
